import os
from datetime import datetime
import baostock as bs
import tushare as ts
import akshare as ak
import json
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.offline as pyo
import numpy as np

from libs.ma import *
from libs.kdj import *
from libs.macd import *

class StockData:
    def __init__(self, code, start_date, end_date, frequency):
        self.code = code
        self.start_date = start_date
        self.end_date = end_date
        self.frequency = frequency
        self.data = None
        self.df = None
        self.filename = f"data\\{self.code}\\{self.code}_{self.start_date}_{self.end_date}_{self.frequency}_kline_data.json"
        self.fig = None
        self.rows = 2
        self.kdj = None
        self.macd = None
        self.k_text = []
        self.k_line = []
    
    def __set_dataframe(self):
        if not self.data:
            print("请先获取股票数据")
            return
        # 将数据转换为DataFrame
        self.df = pd.DataFrame(self.data)
        # 确保日期列是datetime类型
        self.df['date'] = pd.to_datetime(self.df['date'])
        # 确保价格列是float类型
        self.df['open'] = self.df['open'].astype(float).round(2)
        self.df['high'] = self.df['high'].astype(float).round(2)
        self.df['low'] = self.df['low'].astype(float).round(2)
        self.df['close'] = self.df['close'].astype(float).round(2)
        self.df['volume'] = self.df['volume'].astype(float)
        self.df['increase'] = self.df['close'] - self.df['close'].shift(1)
        self.df['pct_increase'] = ((self.df['increase'] / self.df['open']) * 100).round(2)

    def __get_stock_data_from_baostock(self, adjustflag):
        # 登录baostock
        lg = bs.login()
        # 获取数据
        rs = bs.query_history_k_data_plus(self.code,
                                        "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
                                        start_date=self.start_date, end_date=self.end_date,
                                        frequency=self.frequency, adjustflag=adjustflag)  # 3:不复权,2:前复权，1:后复权
        if rs.error_code != '0':
            bs.logout()
            print("获取数据失败，错误码：", rs.error_code)
            return False
        data_list = []
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
        if len(data_list) == 0:
            bs.logout()
            print("baostock 未获取到数据")
            return False
        # 列名
        columns = rs.fields
        # 将数据转换为字典列表
        self.data = [dict(zip(columns, data)) for data in data_list]
        # 登出baostock
        bs.logout()
        # 保存为JSON文件
        if not os.path.exists(os.path.dirname(self.filename)):
            os.makedirs(os.path.dirname(self.filename))
        with open(self.filename, 'w', encoding='utf-8') as f:
            json.dump(self.data, f, ensure_ascii=False, indent=4)
        return True
    
    def __get_stock_data_from_tushare(self):
        # 获取数据
        df = ts.get_k_data(self.code, start=self.start_date, end=self.end_date, ktype=self.frequency)
        if df.empty:
            print("tushare 未获取到数据")
            return False
        # 保存为JSON文件
        if not os.path.exists(os.path.dirname(self.filename)):
            os.makedirs(os.path.dirname(self.filename))
        df.to_json(self.filename, orient='records')
        self.data = df.to_dict(orient='records')
        return True
    
    def __get_stock_data_from_akshare(self):
        # 获取数据
        df = ak.stock_zh_a_hist(symbol=self.code[3:], start_date=datetime.strptime(self.start_date, "%Y-%m-%d").strftime("%Y%m%d"), end_date=datetime.strptime(self.end_date, "%Y-%m-%d").strftime("%Y%m%d"), adjust="qfq")
        if df.empty:
            print("akshare 未获取到数据")
            return False
        # 保存为JSON文件
        if not os.path.exists(os.path.dirname(self.filename)):
            os.makedirs(os.path.dirname(self.filename))
        df.to_json(self.filename, orient='records')
        self.data = df.to_dict(orient='records')
        return True

    def __get_stock_data_from_file(self):
        try:
            with open(self.filename, 'r', encoding='utf-8') as f:
                self.data = json.load(f)
            return True
        except:
            return False
    
    def get_stock_data(self, adjustflag="2"):
        if adjustflag != 2:
            self.filename = self.filename[:-5] + f"_{adjustflag}" + self.filename[-4:]
        ret = self.__get_stock_data_from_file()
        if not ret:
            ret = self.__get_stock_data_from_baostock(adjustflag)
        # if not ret:
        #     ret = self.__get_stock_data_from_tushare()
        #     ret = self.__get_stock_data_from_akshare()
        if not ret:
            print("获取股票数据失败")
            return False
        self.__set_dataframe()
        return self.df
    
    def add_kdj_to_chart(self, n=9, m1=3, m2=3):
        self.kdj = kdj(self.df, n, m1, m2)
        self.rows += 1
    
    def add_macd_to_chart(self, n_fast=12, n_slow=26, n_signal=9):
        self.macd = macd(self.df, n_fast, n_slow, n_signal)
        self.rows += 1

    def __add_kline_to_chart(self, row=1):
        df = self.df
        self.fig.add_trace(go.Candlestick(x=df.index,
                                    open=df['open'],
                                    high=df['high'],
                                    low=df['low'],
                                    close=df['close'], 
                                    increasing_line_color='red', 
                                    decreasing_line_color='green',
                                    hoverinfo='text',
                                    text = df['date'].dt.strftime('%Y-%m-%d') + '<br>' + 
                                            '开盘：' + df['open'].astype(str) + '<br>' + 
                                            '收盘：' + df['close'].astype(str) + '<br>' + 
                                            '最高：' + df['high'].astype(str) + '<br>' + 
                                            '最低：' + df['low'].astype(str) + '<br>' + 
                                            '涨幅：' + df['pct_increase'].astype(str) + '%',
                                    ), 
                                    row=row, col=1)
        # 配置rangeslider
        self.fig.update_xaxes(
            rangeslider_visible=True,  # 显示rangeslider
            rangeslider_thickness=0.05,  # 设置rangeslider的厚度
            range=[df.index[0], df.index[-1]],  # 初始显示范围
            row=row, col=1
        )
        # 配置Y轴动态缩放
        self.fig.update_yaxes(
            autorange=True,  # 允许Y轴自动缩放
            fixedrange=False,  # 取消固定Y轴范围
            row=row, col=1
        )

    def __add_volume_to_chart(self, row=2):
        df = self.df
        # 添加成交量图
        self.fig.add_trace(go.Bar(x=df.index, y=df['volume'], marker=dict(color='blue'), 
                            hoverinfo='text',
                            text = df['date'].dt.strftime('%Y-%m-%d') + '<br>' + 
                            '成交量：' + df['volume'].astype(str) + '<br>'), 
                            row=row, col=1)
    
    def add_text_to_kline_chart(self, i, text, color="#yellow"):
        self.k_text.append((i, text, color))

    def add_line_to_kline_chart(self, x0, y0, x1, y1, color="blue", width=2):
        self.k_line.append((x0, y0, x1, y1, color, width))

    def draw_kline_chart(self):
        rows = 0
        self.fig = make_subplots(rows=self.rows, cols=1, shared_xaxes=True, vertical_spacing=0.02)
        self.__add_kline_to_chart(row=1)
        self.__add_volume_to_chart(row=2)
        rows += 2
        if self.kdj is not None:
            self.fig.add_trace(go.Scatter(x=self.kdj.index, y=self.kdj['k'], name='K值', line=dict(color='blue')), row=rows, col=1)
            self.fig.add_trace(go.Scatter(x=self.kdj.index, y=self.kdj['d'], name='D值', line=dict(color='red')), row=rows, col=1)
            self.fig.add_trace(go.Scatter(x=self.kdj.index, y=self.kdj['j'], name='J值', line=dict(color='green')), row=rows, col=1)
        rows += 1
        if self.macd is not None:
            self.fig.add_trace(go.Scatter(x=self.macd.index, y=self.macd['macd'], name='DIF', line=dict(color='blue')), row=rows, col=1)
            self.fig.add_trace(go.Scatter(x=self.macd.index, y=self.macd['signal'], name='DEA', line=dict(color='red')), row=rows, col=1)
            self.fig.add_trace(go.Bar(x=self.macd.index, y=self.macd['histogram'], name='MACD', marker=dict(color='green')), row=rows, col=1)
        rows += 1

        sma_result = sma(self.df, ma=(5,10,20,30,60,120,250))
        for ma_period in [5, 10, 20, 30, 60, 120, 250]:
            self.fig.add_trace(
                go.Scatter(
                    x=sma_result.index, 
                    y=sma_result[f'ma{ma_period}'], 
                    name=f'MA{ma_period}', 
                    line=dict(width=1),
                    hoverinfo='skip'  # 添加这一行，表示在鼠标悬停时跳过显示
                ), 
                row=1, 
                col=1
            )

        if self.k_text:
            for (i, text, color) in self.k_text:
                self.fig.add_annotation(x=self.df.index[i], y=self.df['high'][i], text=text, showarrow=True, row=1, col=1, bgcolor=color)
        if self.k_line:
            for (x0, y0, x1, y1, color, width) in self.k_line:
                self.fig.add_annotation(x=x1, y=y1, ax=x0, ay=y0, axref='x', ayref='y', xref='x', yref='y',
                showarrow=True, arrowhead=2, arrowsize=1, arrowwidth=width, arrowcolor=color)
        pyo.plot(self.fig, filename=f'{self.filename[:-5]}.html', auto_open=True)
